from multiprocessing import (Manager, Queue, Process, Value, cpu_count, Event)
from sage.combinat.backtrack import SearchForest
from sage.rings.integer import Integer
from sage.misc.abstract_method import abstract_method
from sage.categories.finite_enumerated_sets import FiniteEnumeratedSets
from sage.structure.unique_representation import UniqueRepresentation
import threading
import os

class SearchForest2(SearchForest):
    r""" Commentaire """
    def __init__(self, roots = None, children = None, post_process = None,
        statistic = None, reduce_function = None, reduce_init = None, algorithm = 'depth', category = None ):

        SearchForest.__init__(self,roots,children,post_process, algorithm, category)	
        """ Initialisation du post_process """
        if post_process is None:
            self.post_process = lambda x: x
        if statistic is not None :
            self.statistic = statistic
        if reduce_function is not None :
            self.reduce_function = reduce_function
        self._reduce_init = reduce_init
        self._res = None
        self._manager = Manager()

    @abstract_method
    def statistic(self, x):
        r"""Commentaire"""

    @abstract_method
    def reduce_function(self, x, y):
        r"""Commentaire """

    def traitement(self, queues,fin, ar, i,dor):
        r"""Commentaire """
        qu = queues[i]
        if type(self._reduce_init) == list:
            res = list(self._reduce_init)
        else:
            res = self._reduce_init
        compteur = 0
        while fin.value == 0:
            while qu.qsize() > 0:
                try:
                    o = qu.get(False)	# On lit dans le queue mais pas en bloquant d'ou le False. Si rien a lire exception.
                    for ch in self.children(o): # Ajout des fils dans la queue
                        qu.put(ch)
                    e = self.post_process(o)
                    if e is not None :   # Si post_process return un objet different de none alors
                        res = self.reduce_function(res, self.statistic(e)) # Il enregistre l'objet
                except:
                    continue

            if qu.qsize() > 0:
                continue

            ar.append(i)
            for q in queues:
                if q.qsize() > 0 and id(q) != id(qu):
                    try:
                        element = q.get(timeout = 0.1)
                        ar.remove(i)
                        qu.put(element)
                        break
                    except:
                        continue
            dor.set()

        #print "Processe : ",i," nombre : ",compteur
        if type(self._reduce_init) == list:
            self._res = self._res + res 
        else:
            self._res.value = self.reduce_function(self._res.value,res)

    def run(self):
        r"""Commentaire """
        if self._res is None :
            if type(self._reduce_init) == list:	# Il faut traite le cas des listes.
                self._res = self._manager.list(self._reduce_init)
            else:
                self._res = Value('b',self._reduce_init)
            
            fin = Value('b',0)
            arr = self._manager.list()		# List pour les processuse endormis
            dort = Event()
            queues = [Queue() for i in range(cpu_count())]

            for i in range(len(self.roots())):
                queues[i % len(queues)].put(self.roots()[i])

            processes = [Process(target = self.traitement, args = (queues,fin,arr,i,dort)) for i in range(cpu_count())]

            for p in processes:
                p.start()

            while fin.value == 0:
                dort.wait()
                dort.clear()
                if len(arr) >= cpu_count():
                    print "fin"
                    fin.value = 1

            for p in processes:
                p.join()

        return self._res

    def initRes(self):					# Reinitialisation de res. Afin de refaire le calcul
        if self._res != None :
            self._res = None
